{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "### TAO remote client - Purpose built models\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "![image](https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png)\n",
    "\n",
    "\n",
    "### The workflow in a nutshell\n",
    "\n",
    "- Creating a dataset\n",
    "- Upload dataset to the service\n",
    "- Running dataset convert (for specific models)\n",
    "- Getting a PTM from NGC\n",
    "- Model Actions\n",
    "    - Train (Normal/AutoML)\n",
    "    - Evaluate\n",
    "    - Prune, retrain (for specific models)\n",
    "    - Export\n",
    "    - TAO-Deploy (for specific models)\n",
    "    - Inference on TAO\n",
    "    - Inference on TRT (for specific models)\n",
    "\n",
    "### Table of contents\n",
    "\n",
    "1. [Install TAO remote client ](#head-1)\n",
    "1. [Set the remote service base URL](#head-2)\n",
    "1. [Access the shared volume](#head-3)\n",
    "1. [Create the datasets](#head-4)\n",
    "1. [List datasets](#head-5)\n",
    "1. [Provide and customize dataset convert specs for train dataset](#head-6)\n",
    "1. [Run dataset convert for train dataset](#head-7)\n",
    "1. [Provide and customize dataset convert specs for val dataset](#head-8)\n",
    "1. [Run dataset convert for val dataset](#head-9)\n",
    "1. [Create an experiment](#head-10)\n",
    "1. [Find pretrained model](#head-11)\n",
    "1. [Customize model metadata](#head-12)\n",
    "1. [View hyperparameters that are enabled for AutoML by default](#head-13)\n",
    "1. [Set AutoML related configurations](#head-14)\n",
    "1. [Provide train specs](#head-15)\n",
    "1. [Run train](#head-16)\n",
    "1. [View checkpoint files](#head-17)\n",
    "1. [Provide evaluate specs](#head-18)\n",
    "1. [Run evaluate](#head-19)\n",
    "1. [Provide prune specs](#head-20) (for specific models)\n",
    "1. [Run prune](#head-21) (for specific models)\n",
    "1. [Provide retrain specs](#head-22) (for specific models)\n",
    "1. [Run retrain](#head-23) (for specific models)\n",
    "1. [Run evaluate on retrain](#head-24) (for specific models)\n",
    "1. [Provide export specs](#head-25)\n",
    "1. [Run export](#head-26)\n",
    "1. [Provide trt engine generation specs](#head-27) (for specific models)\n",
    "1. [Run TRT Engine generation using TAO-Deploy](#head-28) (for specific models)\n",
    "1. [Provide TAO inference specs](#head-29)\n",
    "1. [Run TAO inference](#head-30)\n",
    "1. [Provide TRT inference specs](#head-31) (for specific models)\n",
    "1. [Run TRT inference](#head-32) (for specific models)\n",
    "1. [Delete experiment](#head-33)\n",
    "1. [Delete datasets](#head-34)\n",
    "1. [Unmount shared volume](#head-35)\n",
    "1. [Uninstall TAO Remote Client](#head-36)\n",
    "\n",
    "### Requirements\n",
    "Please find the server requirements [here](https://docs.nvidia.com/tao/tao-toolkit/text/tao_toolkit_api/api_setup.html#)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import subprocess\n",
    "import json\n",
    "import time\n",
    "import ast\n",
    "import shutil\n",
    "from IPython.display import clear_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "namespace = 'default'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FIXME\n",
    "\n",
    "1. Assign a model_name in FIXME 1\n",
    "\n",
    "    1.1 Assign model type for action_recognition/pose_classification in FIXME 1.1\n",
    "\n",
    "    1.2 Assign platform for action_recognition in FIXME 1.2\n",
    "    \n",
    "    1.3 Assign model input type for action_recognition in FIXME 1.3\n",
    "1. Assign a workdir in FIXME 2\n",
    "1. Assign the ip_address and port_number in FIXME 3 ([info](https://docs.nvidia.com/tao/tao-toolkit/text/tao_toolkit_api/api_rest_api.html))\n",
    "1. Assign the ngc_api_key variable in FIXME 4\n",
    "1. (Optional) Enable AutoML if needed in FIXME 5\n",
    "1. (Optional) Choose between bayesian and hyperband automl_algorithm in FIXME 6 (If automl was enabled in FIXME5)\n",
    "1. Choose to download jobs or not in FIXME 7\n",
    "1. Choose between default and custom dataset in FIXME 8\n",
    "1. Assign path of DATA_DIR in FIXME 9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Define model_name workspaces and other variables\n",
    "# Available models (#FIXME 1):\n",
    "# 1. action_recognition - https://docs.nvidia.com/tao/tao-toolkit/text/action_recognition_net.html\n",
    "# 2. bpnet - https://docs.nvidia.com/tao/tao-toolkit/text/bodypose_estimation/bodyposenet.html\n",
    "# 3. fpenet - https://docs.nvidia.com/tao/tao-toolkit/text/facial_landmarks_estimation/facial_landmarks_estimation.html\n",
    "# 4. lprnet - https://docs.nvidia.com/tao/tao-toolkit/text/character_recognition/index.html\n",
    "# 5. ml_recog - https://docs.nvidia.com/tao/tao-toolkit/text/ml_recog/index.html\n",
    "# 6. ocdnet - https://docs.nvidia.com/tao/tao-toolkit/text/ocdnet/index.html\n",
    "# 7. ocrnet - https://docs.nvidia.com/tao/tao-toolkit/text/ocrnet/index.html\n",
    "# 8. optical_inspection - https://docs.nvidia.com/tao/tao-toolkit/text/optical_inspection/index.html\n",
    "# 9. pose_classification - https://docs.nvidia.com/tao/tao-toolkit/text/pose_classification/index.html\n",
    "# 10. pointpillars - https://docs.nvidia.com/tao/tao-toolkit/text/point_cloud/pointpillars.html\n",
    "# 11. re_identification - https://docs.nvidia.com/tao/tao-toolkit/text/re_identification/index.html\n",
    "# 12. centerpose\n",
    "# 13. visual_changenet_segment - https://docs.nvidia.com/tao/tao-toolkit/text/visual_changenet/index.html\n",
    "# 14. visual_changenet_classify - https://docs.nvidia.com/tao/tao-toolkit/text/visual_changenet/index.html\n",
    "\n",
    "model_name = \"action_recognition\" # FIXME1 (Add the model name from the above mentioned list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name in (\"action_recognition\",\"fpenet\",\"lprnet\",\"pose_classification\"):\n",
    "    # FIXME1.1 - model_type - string\n",
    "        # action_recognition: rgb/of/joint;\n",
    "        # fpenet: 10/80 (value represents the number of keypoints)\n",
    "        # lprnet: us/ch (us for United States, ch for China)\n",
    "        # pose_classification: kinetics/nvidia\n",
    "    model_type = \"rgb\" # FIXME1.1 action_recognition: rgb/of/joint; fpenet: 10/80 (refers to number of keypoints); pose_classification: kinetics/nvidia\n",
    "\n",
    "    if model_name == \"action_recognition\":\n",
    "        if model_type not in (\"rgb\",\"of\",\"joint\"):\n",
    "            raise Exception(\"Choose one of rgb/of/joint for action recognition model_type\")\n",
    "    elif model_name == \"fpenet\":\n",
    "        if model_type not in (\"10\",\"80\"):\n",
    "            raise Exception(\"Choose one of 10/80 for FPENET model_type\")\n",
    "    elif model_name == \"lprnet\":\n",
    "        if model_type not in (\"us\",\"ch\"):\n",
    "            raise Exception(\"Choose one of us/ch for LPRNET model_type\")\n",
    "    elif model_name == \"pose_classification\":\n",
    "        if model_type not in (\"kinetics\",\"nvidia\"):\n",
    "            raise Exception(\"Choose one of kinetics/nvidia for pose classification model_type\")\n",
    "\n",
    "    if model_name == \"action_recognition\":\n",
    "        platform = \"a100\" # FIXME1.2 a100/xavier - valid only for model_type that is not rgb\n",
    "        model_input_type = \"3d\" # FIXME1.3 3d/2d"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "### Install TAO remote client <a class=\"anchor\" id=\"head-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO-Client wheel.\n",
    "! pip3 install nvidia-transfer-learning-client\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# View the version of the TAO-Client\n",
    "! nvtl --version"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "### Set the remote service base URL and Token <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Define the node_addr and port number\n",
    "workdir = \"workdir_purpose_built_models\" # FIXME2\n",
    "host_url = \"http://<ip_address>:<port_number>\" # FIXME3 example: https://10.137.149.22:32334\n",
    "# In host machine, node ip_address and port number can be obtained as follows,\n",
    "# ip_address: hostname -i\n",
    "# port_number: kubectl get service ingress-nginx-controller -o jsonpath='{.spec.ports[0].nodePort}'\n",
    "\n",
    "ngc_api_key = \"<ngc_api_key>\" # FIXME4 example: (Add NGC API key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "automl_enabled = False # FIXME5 set to True if you want to run automl for the model chosen in the previous cell\n",
    "automl_algorithm=\"bayesian\" # FIXME6 example: bayesian/hyperband\n",
    "# FIXME7 Defaulted to False as downloading jobs from service to your machine takes time\n",
    "# Set to True if you want to download jobs where examples have been provided like for train, export, inference.\n",
    "download_jobs = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%env BASE_URL={host_url}/{namespace}/api/v1\n",
    "\n",
    "# Exchange NGC_API_KEY for JWT\n",
    "identity = json.loads(subprocess.getoutput(f\"nvtl login --ngc-api-key {ngc_api_key}\"))\n",
    "\n",
    "%env USER={identity['user_id']}\n",
    "%env TOKEN={identity['token']}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Creating workdir\n",
    "workdir = os.path.abspath(workdir)\n",
    "if not os.path.isdir(workdir):\n",
    "    os.makedirs(workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function to parse logs <a class=\"anchor\" id=\"head-1.1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def my_tail(model_name_cli, id, job_id, job_type, workdir):\n",
    "\tstatus = None\n",
    "\twhile True:\n",
    "\t\ttime.sleep(10)\n",
    "\t\tclear_output(wait=True)\n",
    "\t\tlog_file_path = subprocess.getoutput(f\"nvtl {model_name_cli} get-log-file --id {id} --job {job_id} --job_type {job_type} --workdir {workdir}\")\n",
    "\t\tif not os.path.exists(log_file_path):\n",
    "\t\t\tcontinue\n",
    "\t\twith open(log_file_path, 'rb') as log_file:\n",
    "\t\t\tlog_contents = log_file.read()\n",
    "\t\tlog_content_lines = log_contents.decode(\"utf-8\").split(\"\\n\")        \n",
    "\t\tfor line in log_content_lines:\n",
    "\t\t\tprint(line.strip())\n",
    "\t\t\tif line.strip() == \"Error EOF\":\n",
    "\t\t\t\tstatus = \"Error\"\n",
    "\t\t\t\tbreak\n",
    "\t\t\telif line.strip() == \"Done EOF\":\n",
    "\t\t\t\tstatus = \"Done\"\n",
    "\t\t\t\tbreak\n",
    "\t\tif status is not None:\n",
    "\t\t\tbreak\n",
    "\treturn status"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function to split tar files <a class=\"anchor\" id=\"head-1.1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tarfile\n",
    "\n",
    "def split_tar_file(input_tar_path, output_dir, max_split_size=0.2*1024*1024*1024):\n",
    "\tos.makedirs(output_dir, exist_ok=True)\n",
    "\t\n",
    "\twith tarfile.open(input_tar_path, 'r') as original_tar:\n",
    "\t\tmembers = original_tar.getmembers()\n",
    "\t\tcurrent_split_size = 0\n",
    "\t\tcurrent_split_number = 0\n",
    "\t\tcurrent_split_name = os.path.join(output_dir, f'smaller_file_{current_split_number}.tar')\n",
    "\t\t\n",
    "\t\twith tarfile.open(current_split_name, 'w') as split_tar:\n",
    "\t\t\tfor member in members:\n",
    "\t\t\t\tif current_split_size + member.size <= max_split_size:\n",
    "\t\t\t\t\tsplit_tar.addfile(member, original_tar.extractfile(member))\n",
    "\t\t\t\t\tcurrent_split_size += member.size\n",
    "\t\t\t\telse:\n",
    "\t\t\t\t\tsplit_tar.close()\n",
    "\t\t\t\t\tcurrent_split_number += 1\n",
    "\t\t\t\t\tcurrent_split_name = os.path.join(output_dir, f'smaller_file_{current_split_number}.tar')\n",
    "\t\t\t\t\tcurrent_split_size = 0\n",
    "\t\t\t\t\tsplit_tar = tarfile.open(current_split_name, 'w')  # Open a new split tar archive\n",
    "\t\t\t\t\tsplit_tar.addfile(member, original_tar.extractfile(member))\n",
    "\t\t\t\t\tcurrent_split_size += member.size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set dataset type, format <a class=\"anchor\" id=\"head-1.1\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "**Action Recognition:** We will be using the HMDB51 [dataset](https://serre-lab.clps.brown.edu/resource/hmdb-a-large-human-motion-database/) for the tutorial. (We choose catch/smile for this tutorial):\n",
    "\n",
    "**BPNET:** We will be using the `COCO dataset` for Instance segmentation - MaskRCNN. `download_coco.sh` script from dataset prepare will be used to download and unzip the coco2017 dataset from [here](https://cocodataset.org/#download)\n",
    "\n",
    "**FPENET:** We will be using `AFW dataset`. Download it from [here](https://ibug.doc.ic.ac.uk/download/annotations/afw.zip/) and place it in $DATA_DIR.\n",
    "\n",
    "**LPRNET**: We will be using the `OpenALPR benchmark dataset` for the tutorial. The following script will download the dataset automatically and convert it to the format used by TAO.  \n",
    "\n",
    "**MLRecogNet** We will be using the `Retail Product Checkout Dataset` for the tutorial. Downdload the datsaet from [here](https://www.kaggle.com/datasets/diyer22/retail-product-checkout-dataset) and place it under $DATA_DIR/metric_learning_recognition\n",
    "\n",
    "**OCDNET**: We will be using the ICDAR2015 dataset for the ocdnet tutorial. Please access the dataset [here](https://rrc.cvc.uab.es/?ch=4&com=tasks) to register and download the data from Task 4.1: Text Localization. Unzip the files to DATA_DIR\n",
    "\n",
    "**OCRNET**: We will be using the ICDAR15 word recognition dataset for the tutorial. To find more details please visit [here](\n",
    "https://rrc.cvc.uab.es/?ch=4&com=tasks). Please download the ICDAR15 word recognition train dataset and test_dataset [here](https://rrc.cvc.uab.es/?ch=4&com=downloads) to DATA_DIR.\n",
    "\n",
    "**Pointpillars:** We will be using the `kitti object detection dataset` for this example. To find more details, please visit [here](http://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=2d)\n",
    "\n",
    "**Pose Classification:** We will be using the Kinetics dataset from [Deepmind](https://deepmind.com/research/open-source/kinetics) or NVIDIA created dataset. For kinetics based dataset set model_type as `kinetics` and for nvidia based dataset set model_type as `nvidia`\n",
    "\n",
    "**re_identification:** We will be using the [Market-1501](https://zheng-lab.cecs.anu.edu.au/Project/project_reid.html) dataset. Download the dataset [here](https://drive.google.com/file/d/1TwkgQcIa_EgRjVMPSbyEKtcfljqURrzi/view?usp=sharing) and extract it.\n",
    "\n",
    "**Optical Inspection:** Bring your own dataset according to the format described [here](https://docs.nvidia.com/tao/tao-toolkit/text/data_annotation_format.html#optical_inspection-format). \n",
    "\n",
    "**Visual ChangeNet-Classification:** Bring your own dataset according to the format described [here](https://docs.nvidia.com/tao/tao-toolkit/text/data_annotation_format.html#optical_inspection-format). \n",
    "\n",
    "**Visual ChangeNet-Segmentation:** We will be using the [Market-1501](https://zheng-lab.cecs.anu.edu.au/Project/project_reid.html) dataset. Download the dataset [here](https://www.dropbox.com/s/18fb5jo0npu5evm/LEVIR-CD256.zip) and extract it. \n",
    "\n",
    "**CenterPose:** We will be using [Google Objectron](https://github.com/google-research-datasets/Objectron) dataset. The following script will download and preprocess the dataset the dataset automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name == \"lprnet\":\n",
    "    ds_type = \"character_recognition\"\n",
    "    ds_format = \"lprnet\"\n",
    "elif model_name in (\"visual_changenet_classify\", \"visual_changenet_segment\"):\n",
    "    ds_type = \"visual_changenet\"\n",
    "    ds_format = model_name\n",
    "    model_name = 'visual_changenet'\n",
    "else:\n",
    "    ds_type = model_name\n",
    "    ds_format = \"default\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "dataset_to_be_used = \"default\" #FIXME8 #default/custom; default for the dataset used in this tutorial notebook; custom for a different dataset\n",
    "DATA_DIR = os.path.abspath(model_name) # FIXME9 (set absolute path of the data_directory)\n",
    "os.environ['DATA_DIR']= DATA_DIR\n",
    "!mkdir -p {DATA_DIR}\n",
    "job_map = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset download and pre-processing <a class=\"anchor\" id=\"head-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if dataset_to_be_used == \"default\":\n",
    "    if model_name == \"action_recognition\":\n",
    "        !sudo apt-get update -y && sudo apt-get install unrar-free -y\n",
    "        !wget -P $DATA_DIR http://serre-lab.clps.brown.edu/wp-content/uploads/2013/10/hmdb51_org.rar\n",
    "        assert os.path.exists(f\"{DATA_DIR}/hmdb51_org.rar\")\n",
    "        !mkdir -p $DATA_DIR/videos && unrar x -o+ $DATA_DIR/hmdb51_org.rar $DATA_DIR/videos\n",
    "        !mkdir -p $DATA_DIR/raw_data\n",
    "        !unrar x -o+ $DATA_DIR/videos/catch.rar $DATA_DIR/raw_data\n",
    "        !unrar x -o+ $DATA_DIR/videos/smile.rar $DATA_DIR/raw_data\n",
    "        assert os.path.exists(f\"{DATA_DIR}/raw_data/catch\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/raw_data/smile\")\n",
    "    elif model_name == \"bpnet\":\n",
    "        !bash ../dataset_prepare/coco/download_coco.sh $DATA_DIR\n",
    "        # Remove existing data\n",
    "        !rm -rf $DATA_DIR/train2017/images\n",
    "        !rm -rf $DATA_DIR/val2017/images\n",
    "        # Rearrange data in the required format\n",
    "        !mv $DATA_DIR/raw-data/* $DATA_DIR/\n",
    "        !cp ../dataset_prepare/bpnet/* $DATA_DIR/\n",
    "        assert os.path.exists(f\"{DATA_DIR}/train2017\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/val2017\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/annotations\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/bpnet_18joints.json\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/coco_spec.json\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/infer_spec.yaml\")\n",
    "    elif model_name == \"fpenet\":\n",
    "        assert os.path.exists(f\"{DATA_DIR}/afw.zip\")\n",
    "        !mkdir $DATA_DIR/data\n",
    "        !unzip -uq $DATA_DIR/afw.zip -d $DATA_DIR/data/afw\n",
    "        !cp ../dataset_prepare/fpenet/data.json $DATA_DIR/\n",
    "        assert os.path.exists(f\"{DATA_DIR}/data/afw\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/data.json\")\n",
    "    elif model_name == \"lprnet\":\n",
    "        !python3 -m pip install --upgrade pip\n",
    "        !python3 -m pip install \"opencv-python>=3.4.0.12,<=4.5.5.64\"\n",
    "        !bash ../dataset_prepare/lprnet/download_and_prepare_data.sh $DATA_DIR\n",
    "        assert os.path.exists(f\"{DATA_DIR}/train/image\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/train/label\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/val/image\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/val/label\")\n",
    "    elif model_name == \"ml_recog\":\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset.zip\")\n",
    "        !unzip -uq $DATA_DIR/metric_learning_recognition/retail-product-checkout-dataset.zip -d $DATA_DIR/metric_learning_recognition\n",
    "    elif model_name == \"ocdnet\":\n",
    "        assert(os.path.exists(f\"{DATA_DIR}/train/img\"))\n",
    "        assert(os.path.exists(f\"{DATA_DIR}/train/gt\"))\n",
    "        assert(os.path.exists(f\"{DATA_DIR}/test/img\"))\n",
    "        assert(os.path.exists(f\"{DATA_DIR}/test/gt\"))\n",
    "    elif model_name == \"ocrnet\":\n",
    "        !mkdir -p $DATA_DIR/train && rm -rf $DATA_DIR/train/*\n",
    "        !mkdir -p $DATA_DIR/test && rm -rf $DATA_DIR/test/*\n",
    "        !unzip -u $DATA_DIR/ch4_test_word_images_gt.zip -d $DATA_DIR/test\n",
    "        !cp $DATA_DIR/Challenge4_Test_Task3_GT.txt -d $DATA_DIR/test\n",
    "        !unzip -u $DATA_DIR/ch4_training_word_images_gt.zip -d $DATA_DIR/train    \n",
    "        assert os.path.exists(f\"{DATA_DIR}/ch4_test_word_images_gt.zip\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/Challenge4_Test_Task3_GT.txt\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/ch4_training_word_images_gt.zip\")\n",
    "    elif model_name == \"optical_inspection\" or (model_name == \"visual_changenet\" and ds_format == \"visual_changenet_classify\"):\n",
    "        assert os.path.exists(f\"{DATA_DIR}/train/images\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/train/dataset.csv\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/val/images\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/val/dataset.csv\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/test/images\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/test/dataset.csv\")\n",
    "    elif model_name == \"visual_changenet\" and ds_format == \"visual_changenet_segment\":\n",
    "        #Download the data\n",
    "        URL_DATASET = \"https://www.dropbox.com/s/18fb5jo0npu5evm/LEVIR-CD256.zip\"\n",
    "        os.environ[\"URL_DATASET\"]=URL_DATASET\n",
    "        !if [ ! -f $DATA_DIR/LEVIR-CD256.zip ]; then wget $URL_DATASET -O $DATA_DIR/LEVIR-CD-256.zip; else echo \"image archive already downloaded\"; fi \n",
    "        # Check the dataset is present\n",
    "        !mkdir -p $DATA_DIR\n",
    "        !if [ ! -f $DATA_DIR/LEVIR-CD-256.zip ]; then echo 'Dataset zip file not found, please download.'; else echo 'Found Dataset zip file.';fi\n",
    "        # unpack \n",
    "        !unzip -u $DATA_DIR/LEVIR-CD-256.zip -d $DATA_DIR\n",
    "        assert os.path.exists(f\"{DATA_DIR}/LEVIR-CD256/A\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/LEVIR-CD256/B\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/LEVIR-CD256/label\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/LEVIR-CD256/list/train.txt\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/LEVIR-CD256/list/val.txt\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/LEVIR-CD256/list/test.txt\")\n",
    "    elif model_name == \"pointpillars\":\n",
    "        !unzip -u $DATA_DIR/data_object_image_2.zip -d $DATA_DIR\n",
    "        !unzip -u $DATA_DIR/data_object_label_2.zip -d $DATA_DIR\n",
    "        !unzip -u $DATA_DIR/data_object_velodyne.zip -d $DATA_DIR\n",
    "        !unzip -u $DATA_DIR/data_object_calib.zip -d $DATA_DIR\n",
    "        assert os.path.exists(f\"{DATA_DIR}/training/image_2\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/training/label_2\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/training/velodyne\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/training/calib\")\n",
    "    elif model_name == \"pose_classification\":\n",
    "        !pip3 install -U gdown\n",
    "        if model_type == \"kinetics\":\n",
    "            !gdown https://drive.google.com/uc?id=1dmzCRQsFXJ18BlXj1G9sbDnsclXIdDdR -O $DATA_DIR/st-gcn-processed-data.zip\n",
    "            !unzip $DATA_DIR/st-gcn-processed-data.zip -d $DATA_DIR\n",
    "            !mv $DATA_DIR/data/Kinetics/kinetics-skeleton $DATA_DIR/kinetics\n",
    "            !rm -r $DATA_DIR/data\n",
    "            !rm $DATA_DIR/st-gcn-processed-data.zip\n",
    "            assert os.path.exists(f\"{DATA_DIR}/kinetics\")\n",
    "        elif model_type == \"nvidia\":\n",
    "            !gdown https://drive.google.com/uc?id=1GhSt53-7MlFfauEZ2YkuzOaZVNIGo_c- -O $DATA_DIR/data_3dbp_nvidia.zip\n",
    "            !mkdir -p $DATA_DIR/nvidia\n",
    "            !unzip $DATA_DIR/data_3dbp_nvidia.zip -d $DATA_DIR/nvidia\n",
    "            !rm $DATA_DIR/data_3dbp_nvidia.zip\n",
    "            assert os.path.exists(f\"{DATA_DIR}/nvidia\")\n",
    "            assert os.path.exists(f\"{DATA_DIR}/{model_type}/train_data.npy\") and os.path.exists(f\"{DATA_DIR}/{model_type}/train_label.pkl\") and os.path.exists(f\"{DATA_DIR}/{model_type}/val_data.npy\") and os.path.exists(f\"{DATA_DIR}/{model_type}/val_label.pkl\")\n",
    "    elif model_name == \"re_identification\":\n",
    "        !pip3 install -U gdown\n",
    "        !gdown https://drive.google.com/uc?id=0B8-rUzbwVRk0c054eEozWG9COHM -O $DATA_DIR/market1501.zip\n",
    "        !unzip -u $DATA_DIR/market1501.zip -d $DATA_DIR\n",
    "        !rm -rf $DATA_DIR/market1501\n",
    "        !mv $DATA_DIR/Market-1501-v15.09.15 $DATA_DIR/market1501\n",
    "        !rm $DATA_DIR/market1501.zip\n",
    "        assert os.path.exists(f\"{DATA_DIR}/market1501\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name in (\"lprnet\",\"ocdnet\",\"ocrnet\", \"optical_inspection\") or ds_format == \"visual_changenet_classify\":\n",
    "    eval_dataset_path = f\"{DATA_DIR}/purpose_built_models_val.tar.gz\"\n",
    "if model_name in (\"lprnet\", \"optical_inspection\") or ds_format == \"visual_changenet_classify\":\n",
    "    test_dataset_path = f\"{DATA_DIR}/purpose_built_models_test.tar.gz\"\n",
    "train_dataset_path = f\"{DATA_DIR}/purpose_built_models_train.tar.gz\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "train_dataset_id = subprocess.getoutput(f\"nvtl {model_name} dataset-create --dataset_type {ds_type} --dataset_format {ds_format}\")\n",
    "print(train_dataset_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "if dataset_to_be_used == \"default\":\n",
    "    USER_EXPERIMENT_DIR = os.path.join(\"/shared/users\",os.environ[\"USER\"],\"datasets\",train_dataset_id)\n",
    "    if model_name == \"action_recognition\":\n",
    "        !python3 -m pip install opencv-python numpy\n",
    "        # For rgb action recognition\n",
    "        !if [ -d tao_toolkit_recipes ]; then rm -rf tao_toolkit_recipes; fi\n",
    "        !git clone https://github.com/NVIDIA-AI-IOT/tao_toolkit_recipes\n",
    "        assert os.path.exists(\"tao_toolkit_recipes\")\n",
    "        !cd tao_toolkit_recipes/tao_action_recognition/data_generation/ && bash ./preprocess_HMDB_RGB.sh $DATA_DIR/raw_data $DATA_DIR/processed_data \n",
    "\n",
    "        # For optical flow, comment the above 3 lines and uncomment the below (Note: for generating optical flow, a Turing or Ampere above GPU is needed.)\n",
    "        #!echo <passwd> | sudo -S apt install -y libfreeimage-dev\n",
    "        #!cp ../dataset_prepare/action_recognition/AppOFCuda tao_toolkit_recipes/tao_action_recognition/data_generation/\n",
    "        #!cd tao_toolkit_recipes/tao_action_recognition/data_generation/ && bash ./preprocess_HMDB.sh $DATA_DIR/raw_data $DATA_DIR/processed_data\n",
    "\n",
    "        # download the split files and unrar\n",
    "        !wget -P $DATA_DIR http://serre-lab.clps.brown.edu/wp-content/uploads/2013/10/test_train_splits.rar\n",
    "        assert os.path.exists(f\"{DATA_DIR}/test_train_splits.rar\")\n",
    "        !mkdir -p $DATA_DIR/splits && unrar x -o+ $DATA_DIR/test_train_splits.rar $DATA_DIR/splits\n",
    "        assert os.path.exists(f\"{DATA_DIR}/splits\")\n",
    "        # run split_HMDB to generate training split\n",
    "        !if [ -d $DATA_DIR/train ]; then rm -rf $DATA_DIR/train $DATA_DIR/test; fi\n",
    "        !cd tao_toolkit_recipes/tao_action_recognition/data_generation/ && python3 ./split_dataset.py $DATA_DIR/processed_data $DATA_DIR/splits/testTrainMulti_7030_splits $DATA_DIR/train  $DATA_DIR/test\n",
    "        assert os.path.exists(f'{DATA_DIR}/train')\n",
    "        assert os.path.exists(f'{DATA_DIR}/test')\n",
    "\n",
    "        if os.path.exists(\"tao_toolkit_recipes\"):\n",
    "            shutil.rmtree(\"tao_toolkit_recipes\")\n",
    "\n",
    "        assert not os.path.exists(\"tao_toolkit_recipes\")\n",
    "\n",
    "    elif model_name == \"fpenet\":\n",
    "        !pip3 install numpy opencv-python\n",
    "        if model_type == \"80\":\n",
    "            output_json_path = os.path.join(os.environ['DATA_DIR'], 'data/afw/afw.json')\n",
    "        elif model_type == \"10\":\n",
    "            output_json_path = os.path.join(os.environ['DATA_DIR'], 'data/afw_10/afw_10.json')\n",
    "        !python3 ../dataset_prepare/fpenet/data_utils.py --afw_data_path $DATA_DIR/data/afw --output_json_path $output_json_path --afw_image_save_path $DATA_DIR/data/afw --num_key_points $model_type --container_root_path $USER_EXPERIMENT_DIR\n",
    "        assert os.path.exists(output_json_path)\n",
    "        with open(output_json_path, encoding='utf-8') as afw_json_file:\n",
    "            afw_json = json.load(afw_json_file)\n",
    "            assert afw_json\n",
    "\n",
    "    elif model_name == \"lprnet\":\n",
    "        character_file_link = \"https://api.ngc.nvidia.com/v2/models/nvidia/tao/lprnet/versions/trainable_v1.0/files/{}_lp_characters.txt\".format(model_type)\n",
    "        !wget -q -O $DATA_DIR/train/characters.txt $character_file_link\n",
    "        !cp $DATA_DIR/train/characters.txt $DATA_DIR/val/characters.txt\n",
    "        assert os.path.exists(f\"{DATA_DIR}/train/characters.txt\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/val/characters.txt\")\n",
    "\n",
    "    elif model_name == \"ml_recog\":\n",
    "        # crops images from detection set and form a classification set\n",
    "        # splits to reference/train/val/test set\n",
    "        !sudo apt-get update && sudo apt-get install gcc -y\n",
    "        !python3 -m pip install opencv-python numpy pycocotools tqdm\n",
    "        !python3 ../dataset_prepare/metric_learning_recognition/process_retail_product_checkout_dataset.py\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes/train\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes/train\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes/test\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes/test\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes/val\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes/val\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes/reference\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes/reference\")\n",
    "\n",
    "    elif model_name == \"ocrnet\":\n",
    "        orig_train_gt_file=os.path.join(os.getenv(\"DATA_DIR\"), \"train\", \"gt.txt\")\n",
    "        processed_train_gt_file=os.path.join(os.getenv(\"DATA_DIR\"), \"train\", \"gt_new.txt\")\n",
    "        orig_test_gt_file=os.path.join(os.getenv(\"DATA_DIR\"), \"test\", \"Challenge4_Test_Task3_GT.txt\")\n",
    "        processed_test_gt_file=os.path.join(os.getenv(\"DATA_DIR\"), \"test\", \"gt_new.txt\")\n",
    "        !python3 ../dataset_prepare/ocrnet/preprocess_label.py $orig_train_gt_file $processed_train_gt_file\n",
    "        !python3 ../dataset_prepare/ocrnet/preprocess_label.py $orig_test_gt_file $processed_test_gt_file\n",
    "\n",
    "    elif model_name == \"pointpillars\":\n",
    "        !python3 -m pip install scikit-image numpy\n",
    "        !mkdir -p $DATA_DIR/train/lidar $DATA_DIR/train/label $DATA_DIR/val/lidar $DATA_DIR/val/label\n",
    "\n",
    "        # Convert labels from Camera coordinate system to LIDAR coordinate system, etc\n",
    "        !python3 ../dataset_prepare/pointpillars/gen_lidar_points.py -p $DATA_DIR/training/velodyne \\\n",
    "                                               -c $DATA_DIR/training/calib    \\\n",
    "                                               -i $DATA_DIR/training/image_2  \\\n",
    "                                               -o $DATA_DIR/train/lidar\n",
    "        assert os.listdir(f\"{DATA_DIR}/train/lidar\")\n",
    "        # Drop DontCare class\n",
    "        !python3 ../dataset_prepare/pointpillars/gen_lidar_labels.py -l $DATA_DIR/training/label_2 \\\n",
    "                                               -c $DATA_DIR/training/calib \\\n",
    "                                               -o $DATA_DIR/train/label\n",
    "        # train/val split\n",
    "        !python3 ../dataset_prepare/pointpillars/drop_class.py $DATA_DIR/train/label DontCare\n",
    "        assert os.listdir(f\"{DATA_DIR}/train/label\")\n",
    "        # Change the val set id's if you need a different set of validation images\n",
    "        !python3 ../dataset_prepare/pointpillars/kitti_split.py ../dataset_prepare/pointpillars/val.txt \\\n",
    "                                          $DATA_DIR/train/lidar \\\n",
    "                                          $DATA_DIR/train/label \\\n",
    "                                          $DATA_DIR/val/lidar \\\n",
    "                                          $DATA_DIR/val/label\n",
    "        assert os.listdir(f\"{DATA_DIR}/val/label\")\n",
    "        assert os.listdir(f\"{DATA_DIR}/val/lidar\")\n",
    "\n",
    "    elif model_name == \"pose_classification\" and model_type == \"kinetics\":\n",
    "        !pip3 install numpy\n",
    "        # select actions\n",
    "        !python3 ../dataset_prepare/pose_classification/select_subset_actions.py\n",
    "        assert os.path.exists(f\"{DATA_DIR}/{model_type}/train_data.npy\") and os.path.exists(f\"{DATA_DIR}/{model_type}/train_label.pkl\") and os.path.exists(f\"{DATA_DIR}/{model_type}/val_data.npy\") and os.path.exists(f\"{DATA_DIR}/{model_type}/val_label.pkl\")\n",
    "\n",
    "    elif model_name == \"re_identification\":\n",
    "        #100 is the number of samples to be present in the subset data - you can choose any number <= total samples in the dataset\n",
    "        !python3 ../dataset_prepare/re_identification/obtain_subset_data.py 100\n",
    "        assert os.path.exists(f\"{DATA_DIR}/market1501/sample_train\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/market1501/sample_test\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/market1501/sample_query\")\n",
    "\n",
    "    elif model_name == \"centerpose\":\n",
    "        # Select the training categories: bike, book, bottle, camera, cereal_box, chair, laptop, shoe\n",
    "        # Please set the \"n\" to -1 if you want to run the whole dataset training.\n",
    "        testing_categories = 'bike'\n",
    "        !pip3 install numpy opencv-python tqdm scipy==1.9.2 tensorflow==2.14.0\n",
    "        !python3 ../dataset_prepare/centerpose/prepare_centerpose_dataset.py \\\n",
    "                                            -c $testing_categories \\\n",
    "                                            -n 100\n",
    "        assert os.path.exists(f\"{DATA_DIR}/{testing_categories}/train\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/{testing_categories}/test\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/{testing_categories}/val\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tar the datasets <a class=\"anchor\" id=\"head-1.3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name == \"action_recognition\":\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_train.tar.gz train test\n",
    "elif model_name == \"bpnet\":\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_train.tar.gz train2017 val2017 annotations bpnet_18joints.json  coco_spec.json  infer_spec.yaml\n",
    "elif model_name == \"fpenet\":\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_train.tar.gz data data.json\n",
    "elif model_name == \"lprnet\":\n",
    "    !tar -C $DATA_DIR/train/ -czf $DATA_DIR/purpose_built_models_train.tar.gz image label characters.txt\n",
    "    !tar -C $DATA_DIR/val/ -czf $DATA_DIR/purpose_built_models_val.tar.gz image label characters.txt\n",
    "    !tar -C $DATA_DIR/val/ -czf $DATA_DIR/purpose_built_models_test.tar.gz image\n",
    "elif model_name == \"ml_recog\":\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_train.tar.gz metric_learning_recognition/retail-product-checkout-dataset_classification_demo/\n",
    "elif model_name == \"ocdnet\":\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_train.tar.gz train\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_val.tar.gz test\n",
    "elif model_name == \"ocrnet\":\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_train.tar.gz train character_list\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_val.tar.gz test character_list\n",
    "elif model_name == \"optical_inspection\" or ds_format == \"visual_changenet_classify\":\n",
    "    !tar -C $DATA_DIR/train -czf $DATA_DIR/purpose_built_models_train.tar.gz images dataset.csv\n",
    "    !tar -C $DATA_DIR/val -czf $DATA_DIR/purpose_built_models_val.tar.gz images dataset.csv\n",
    "    !tar -C $DATA_DIR/test -czf $DATA_DIR/purpose_built_models_test.tar.gz images dataset.csv\n",
    "elif ds_format == \"visual_changenet_segment\":\n",
    "    !tar -C $DATA_DIR/LEVIR-CD256 -czf $DATA_DIR/purpose_built_models_train.tar.gz A B list label\n",
    "elif model_name == \"pointpillars\":\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_train.tar.gz train val\n",
    "elif model_name == \"pose_classification\":\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_train.tar.gz $model_type\n",
    "elif model_name == \"re_identification\":\n",
    "    !tar -C $DATA_DIR/market1501 -czf $DATA_DIR/purpose_built_models_train.tar.gz sample_train sample_test sample_query\n",
    "elif model_name == 'centerpose':\n",
    "    !tar -C $DATA_DIR/{testing_categories} -czf $DATA_DIR/purpose_built_models_train.tar.gz train val test"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Upload train dataset <a class=\"anchor\" id=\"head-1.3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "output_dir = os.path.join(os.path.dirname(os.path.abspath(train_dataset_path)), model_name, \"train\")\n",
    "split_tar_file(train_dataset_path, output_dir)\n",
    "for idx, tar_dataset_path in enumerate(os.listdir(output_dir)):\n",
    "    print(f\"Uploading {idx+1}/{len(os.listdir(output_dir))} tar split\")\n",
    "    upload_train_dataset_message = subprocess.getoutput(f\"nvtl {model_name} dataset-upload --id {train_dataset_id} --path {os.path.join(output_dir,tar_dataset_path)}\")\n",
    "    print(upload_train_dataset_message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create and upload val dataset <a class=\"anchor\" id=\"head-1.3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name in (\"lprnet\",\"ocdnet\",\"ocrnet\",\"optical_inspection\") or ds_format == \"visual_changenet_classify\":\n",
    "    eval_dataset_id = subprocess.getoutput(f\"nvtl {model_name} dataset-create --dataset_type {ds_type} --dataset_format {ds_format}\")\n",
    "    print(eval_dataset_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name in (\"lprnet\",\"ocdnet\",\"ocrnet\",\"optical_inspection\") or ds_format == \"visual_changenet_classify\":\n",
    "    output_dir = os.path.join(os.path.dirname(os.path.abspath(eval_dataset_path)), model_name, \"val\")\n",
    "    split_tar_file(eval_dataset_path, output_dir)\n",
    "    for idx, tar_dataset_path in enumerate(os.listdir(output_dir)):\n",
    "        print(f\"Uploading {idx+1}/{len(os.listdir(output_dir))} tar split\")\n",
    "        upload_val_dataset_message = subprocess.getoutput(f\"nvtl {model_name} dataset-upload --id {eval_dataset_id} --path {os.path.join(output_dir,tar_dataset_path)}\")\n",
    "        print(upload_val_dataset_message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create and upload test dataset <a class=\"anchor\" id=\"head-1.4\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name in (\"lprnet\", \"optical_inspection\") or ds_format == \"visual_changenet_classify\":\n",
    "    if model_name == \"lprnet\":\n",
    "        ds_type = \"character_recognition\"\n",
    "        ds_format = \"raw\"\n",
    "    elif 'visual_changenet' in ds_type:\n",
    "        ds_type = ds_type\n",
    "        ds_format = ds_format\n",
    "    else:\n",
    "        ds_type = model_name\n",
    "        ds_format = \"default\"\n",
    "    test_dataset_id = subprocess.getoutput(f\"nvtl {model_name} dataset-create --dataset_type {ds_type} --dataset_format {ds_format}\")\n",
    "    print(test_dataset_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name in (\"lprnet\", \"optical_inspection\") or ds_format == \"visual_changenet_classify\":\n",
    "    output_dir = os.path.join(os.path.dirname(os.path.abspath(test_dataset_path)), model_name, \"test\")\n",
    "    split_tar_file(test_dataset_path, output_dir)\n",
    "    for idx, tar_dataset_path in enumerate(os.listdir(output_dir)):\n",
    "        print(f\"Uploading {idx+1}/{len(os.listdir(output_dir))} tar split\")\n",
    "        upload_test_dataset_message = subprocess.getoutput(f\"nvtl {model_name} dataset-upload --id {test_dataset_id} --path {os.path.join(output_dir,tar_dataset_path)}\")\n",
    "        print(upload_test_dataset_message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### List the created datasets <a class=\"anchor\" id=\"head-5\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "message = subprocess.getoutput(f\"nvtl {model_name} list-datasets\")\n",
    "message = ast.literal_eval(message)\n",
    "for rsp in message:\n",
    "    rsp_keys = rsp.keys()\n",
    "    assert \"id\" in rsp_keys\n",
    "    assert \"type\" in rsp_keys\n",
    "    assert \"format\" in rsp_keys\n",
    "    assert \"name\" in rsp_keys\n",
    "    print(rsp[\"id\"],\"\\t\",rsp[\"type\"],\"\\t\",rsp[\"format\"],\"\\t\\t\",rsp[\"name\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train Dataset convert Action <a class=\"anchor\" id=\"head-3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default train dataset specs\n",
    "if model_name in (\"bpnet\", \"fpenet\", \"ocrnet\", \"pointpillars\"):\n",
    "    train_ds_convert_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --action dataset_convert --id {train_dataset_id} --job_type dataset\")\n",
    "    train_ds_convert_specs = json.loads(train_ds_convert_specs)\n",
    "    print(json.dumps(train_ds_convert_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize train dataset specs\n",
    "if model_name in (\"bpnet\", \"fpenet\", \"ocrnet\", \"pointpillars\"):\n",
    "    if model_name == \"bpnet\":\n",
    "        train_ds_convert_specs[\"mode\"] = \"train\"\n",
    "    elif model_name == \"fpenet\":\n",
    "        train_ds_convert_specs[\"num_keypoints\"] = int(model_type)\n",
    "    print(json.dumps(train_ds_convert_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name in (\"bpnet\", \"fpenet\", \"ocrnet\", \"pointpillars\"):\n",
    "    job_id = subprocess.getoutput(f\"nvtl {model_name} dataset-run-action --id {train_dataset_id}  --action dataset_convert --specs '{json.dumps(train_ds_convert_specs)}'\")\n",
    "    job_map[\"train_convert_\" + model_name] = job_id\n",
    "    print(job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "if model_name in (\"bpnet\", \"fpenet\", \"ocrnet\", \"pointpillars\"):\n",
    "    status = my_tail(model_name, train_dataset_id, job_id, \"dataset\", workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Eval Dataset convert Action <a class=\"anchor\" id=\"head-3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default val dataset specs\n",
    "if model_name == \"bpnet\":\n",
    "    eval_ds_convert_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --action dataset_convert --id {train_dataset_id} --job_type dataset\")\n",
    "elif model_name == \"ocrnet\":\n",
    "    eval_ds_convert_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --action dataset_convert --id {eval_dataset_id} --job_type dataset\")\n",
    "\n",
    "if model_name in (\"bpnet\", \"ocrnet\"):\n",
    "    eval_ds_convert_specs = json.loads(eval_ds_convert_specs)\n",
    "    print(json.dumps(eval_ds_convert_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize val dataset specs\n",
    "if model_name in (\"bpnet\", \"ocrnet\"):\n",
    "    # Apply changes to the specs dictionary if necessary\n",
    "    if model_name == \"bpnet\":\n",
    "        eval_ds_convert_specs[\"mode\"] = \"test\"\n",
    "    print(json.dumps(eval_ds_convert_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name in (\"bpnet\", \"ocrnet\"):\n",
    "    train_convert_job_id = job_map[\"train_convert_\" + model_name]\n",
    "    if model_name == \"bpnet\":\n",
    "        job_id = subprocess.getoutput(f\"nvtl {model_name} dataset-run-action --id {train_dataset_id} --action dataset_convert --parent_job_id {train_convert_job_id} --specs '{json.dumps(eval_ds_convert_specs)}'\")\n",
    "    elif model_name == \"ocrnet\":\n",
    "        job_id = subprocess.getoutput(f\"nvtl {model_name} dataset-run-action --id {eval_dataset_id} --action dataset_convert --parent_job_id {train_convert_job_id} --specs '{json.dumps(eval_ds_convert_specs)}'\")\n",
    "    job_map[\"eval_convert_\" + model_name] = job_id\n",
    "    print(job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "if model_name in (\"bpnet\", \"ocrnet\"):\n",
    "    if model_name == \"bpnet\":\n",
    "        status = my_tail(model_name, train_dataset_id, job_id, \"dataset\", workdir)\n",
    "    elif model_name == \"ocrnet\":\n",
    "        status = my_tail(model_name, eval_dataset_id, job_id, \"dataset\", workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create an experiment <a class=\"anchor\" id=\"head-8\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name in (\"action_recognition\", \"pose_classification\", \"ml_recog\", \"ocrnet\", \"ocdnet\", \"optical_inspection\", \"re_identification\"):\n",
    "    encode_key = \"nvidia_tao\"\n",
    "elif model_name == \"pointpillars\":\n",
    "    encode_key = \"tlt_encode\"\n",
    "else:\n",
    "    encode_key = \"nvidia_tlt\"\n",
    "\n",
    "network_arch = model_name\n",
    "experiment_id = subprocess.getoutput(f\"nvtl {model_name} experiment-create --network_arch {network_arch} --encryption_key {encode_key} \")\n",
    "print(experiment_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Assign train, eval datasets <a class=\"anchor\" id=\"head-10\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "docker_env_vars = {} # Update any variables to be included while triggering Docker run-time like MLOPs variables \n",
    "dataset_information = {}\n",
    "dataset_information[\"train_datasets\"] = [train_dataset_id]\n",
    "if model_name in (\"bpnet\",\"fpenet\",\"lprnet\",\"ml_recog\",\"ocdnet\",\"ocrnet\"):\n",
    "    dataset_information[\"calibration_dataset\"] = train_dataset_id\n",
    "if model_name in (\"lprnet\", \"ocdnet\", \"ocrnet\", \"optical_inspection\") or ds_format == \"visual_changenet_classify\":\n",
    "    dataset_information[\"eval_dataset\"] = eval_dataset_id\n",
    "if model_name in (\"lprnet\", \"optical_inspection\") or ds_format == \"visual_changenet_classify\":\n",
    "    dataset_information[\"inference_dataset\"] = test_dataset_id\n",
    "if model_name in (\"centerpose\"):\n",
    "    dataset_information[\"eval_dataset\"] = train_dataset_id\n",
    "    dataset_information[\"inference_dataset\"] = train_dataset_id\n",
    "dataset_information[\"docker_env_vars\"] = docker_env_vars\n",
    "patched_model = subprocess.getoutput(f\"nvtl {model_name} patch-artifact-metadata --id {experiment_id} --job_type experiment --update_info '{json.dumps(dataset_information)}' \")\n",
    "print(patched_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### List experiments <a class=\"anchor\" id=\"head-5\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# List all pretrained models for the chosen network architecture\n",
    "filter_params = {\"network_arch\": network_arch}\n",
    "message = subprocess.getoutput(f\"nvtl {model_name} list-experiments --filter_params '{json.dumps(filter_params)}'\")\n",
    "message = ast.literal_eval(message)\n",
    "for rsp in message:\n",
    "    rsp_keys = rsp.keys()\n",
    "    if \"encryption_key\" not in rsp.keys():\n",
    "        assert \"name\" in rsp_keys and \"version\" in rsp_keys and \"ngc_path\" in rsp_keys and \"additional_id_info\" in rsp_keys\n",
    "        print(f'PTM Name: {rsp[\"name\"]}; PTM version: {rsp[\"version\"]}; NGC PATH: {rsp[\"ngc_path\"]}; Additional info: {rsp[\"additional_id_info\"]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Assign PTM <a class=\"anchor\" id=\"head-7\"></a>\n",
    "\n",
    "Search for PTM on NGC for the Segmentation model chosen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Assigning pretrained models to different purpose built models versions\n",
    "# From the output of previous cell make the appropriate changes to this map if you want to change the default PTM backbone.\n",
    "# Changing the default backbone here requires changing default spec/config during train/eval etc like for example\n",
    "# If you are changing the ptm to resnet34, then you have to modify the config key num_layers if it exists to 34 manually\n",
    "visual_changenet_ptm = \"visual_changenet_segmentation_levircd:visual_changenet_levircd_trainable_v1.0\" # For segmentation\n",
    "if model_name == \"visual_changenet\" and ds_format == 'visual_changenet_classify':\n",
    "    visual_changenet_ptm = \"visual_changenet_classification:visual_changenet_nvpcb_trainable_v1.0\"\n",
    "pretrained_map = {\"action_recognition\":\"actionrecognitionnet:trainable_v1.0\",\n",
    "                  \"bpnet\" : \"bodyposenet:trainable_v1.0\",\n",
    "                  \"fpenet\" : \"fpenet:trainable_v1.0\",\n",
    "                  \"lprnet\": \"lprnet:trainable_v1.0\",\n",
    "                  \"ml_recog\": \"retail_object_recognition:trainable_v1.0\",\n",
    "                  \"ocdnet\": \"ocdnet:trainable_resnet18_v1.0\",\n",
    "                  \"ocrnet\": \"ocrnet:trainable_v1.0\",\n",
    "                  \"optical_inspection\": \"optical_inspection:trainable_v1.0\",\n",
    "                  \"pointpillars\":\"pointpillarnet:trainable_v1.0\",\n",
    "                  \"pose_classification\":\"poseclassificationnet:trainable_v1.0\",\n",
    "                  \"re_identification\":\"reidentificationnet:trainable_v1.1\",\n",
    "                  \"visual_changenet\":visual_changenet_ptm,\n",
    "                  \"centerpose\": \"pretrained_fan_classification_nvimagenet:fan_small_hybrid_nvimagenet\"}\n",
    "\n",
    "if model_name == \"action_recognition\":\n",
    "    if model_type == \"of\":\n",
    "        pretrained_map[\"action_recognition\"] = \"actionrecognitionnet:trainable_v2.0\"\n",
    "    elif model_type == \"joint\":\n",
    "        pretrained_map[\"action_recognition\"] = \"actionrecognitionnet:trainable_v1.0,actionrecognitionnet:trainable_v2.0\"\n",
    "\n",
    "no_ptm_models = set([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if network_arch not in no_ptm_models:\n",
    "    filter_params = {\"network_arch\": network_arch}\n",
    "    message = subprocess.getoutput(f\"nvtl {model_name} list-experiments --filter_params '{json.dumps(filter_params)}'\")\n",
    "    message = ast.literal_eval(message)\n",
    "\n",
    "    ptm_model_names = pretrained_map[network_arch].split(\",\")\n",
    "    ptm = []\n",
    "\n",
    "    # Search for ptm with given ngc path\n",
    "    for ptm_model_name in ptm_model_names:\n",
    "        ptm_id = None\n",
    "        for rsp in message:\n",
    "            rsp_keys = rsp.keys()\n",
    "            assert \"ngc_path\" in rsp_keys\n",
    "            if rsp[\"ngc_path\"].endswith(pretrained_map[network_arch]):\n",
    "                additional_id_info = []\n",
    "                if rsp[\"additional_id_info\"]:\n",
    "                    assert \"additional_id_info\" in rsp_keys\n",
    "                    additional_id_info = rsp[\"additional_id_info\"].split(\",\")\n",
    "                if (len(additional_id_info) == 0) or \\\n",
    "                    (network_arch == \"lprnet\" and len(additional_id_info) == 1 and additional_id_info[0] == model_type) or \\\n",
    "                    (network_arch == \"action_recognition\" and len(additional_id_info) == 1 and additional_id_info[0] == model_input_type) or \\\n",
    "                    (network_arch == \"action_recognition\" and len(additional_id_info) == 2 and additional_id_info[0] == platform and additional_id_info[1] == model_input_type):\n",
    "                    assert \"id\" in rsp_keys\n",
    "                    ptm_id = rsp[\"id\"]\n",
    "                    print(\"Metadata for model with requested NGC Path\")\n",
    "                    print(rsp)\n",
    "                    break\n",
    "        ptm.append(ptm_id)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if network_arch not in no_ptm_models:\n",
    "    ptm_information = {\"base_experiment\":ptm}\n",
    "    patched_model = subprocess.getoutput(f\"nvtl {model_name} patch-artifact-metadata --id {experiment_id} --job_type experiment --update_info '{json.dumps(ptm_information)}' \")\n",
    "    print(patched_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### View hyperparameters that are enabled for AutoML by default <a class=\"anchor\" id=\"head-13\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if automl_enabled:\n",
    "    # View default automl specs enabled\n",
    "    ! nvtl {model_name} model-automl-defaults --id {experiment_id}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train <a class=\"anchor\" id=\"head-11\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "#### Set AutoML related configurations <a class=\"anchor\" id=\"head-14\"></a>\n",
    "Refer to these hyper-links to see the parameters supported by each network and add more parameters if necessary in addition to the default automl enabled parameters:\n",
    "\n",
    "[ActionRecognitionNet](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/action_recognition/action_recognition%20-%20train.csv), \n",
    "[BPNET](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/bpet/bpnet%20-%20train.csv), \n",
    "[FPENET](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/fpenet/fpenet%20-%20train.csv), \n",
    "[LPRNET](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/lprnet/lprnet%20-%20train.csv), \n",
    "[MetricLearningRecognition](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/ml_recog/ml_recog%20-%20train.csv), \n",
    "[OCDNET](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/ocdnet/ocdnet%20-%20train.csv), \n",
    "[OCRNET](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/ocrnet/ocrnet%20-%20train.csv), \n",
    "[OpticalInspection](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/optical_inspection/optical_inspection%20-%20train.csv), \n",
    "[Pointpillars](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/pointpillars/pointpillars%20-%20train.csv), \n",
    "[PoseClassificationNet](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/pose_classification/pose_classification%20-%20train.csv), \n",
    "[ReIdentificationNet](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/re_identification/re_identification%20-%20train.csv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if automl_enabled:\n",
    "    # Choose any metric that is present in the kpi dictionary present in the model's status.json. \n",
    "    # Example status.json for each model can be found in the respective section in NVIDIA TAO DOCS here: https://docs.nvidia.com/tao/tao-toolkit/text/model_zoo/cv_models/index.html\n",
    "    metric=\"kpi\"\n",
    "\n",
    "    additional_automl_parameters = [] #Refer to parameter list mentioned in the above links and add any extra parameter in addition to the default enabled ones\n",
    "    remove_default_automl_parameters = [] #Remove any hyperparameters that are enabled by default for AutoML\n",
    "\n",
    "    automl_information = {\"automl_enabled\":automl_enabled,\n",
    "                          \"automl_algorithm\":automl_algorithm,\n",
    "                          \"automl_max_recommendations\": 20, # Only for bayesian\n",
    "                          \"automl_R\": 27, # Only for hyperband\n",
    "                          \"automl_nu\": 3, # Only for hyperband\n",
    "                          \"epoch_multiplier\": 1, # Only for hyperband\n",
    "                          # Enable this if you want to add parameters to automl_add_hyperparameters below that are disabled by TAO in the automl_enabled column of the spec csv.\n",
    "                          # Warning: The parameters that are disabled are not tested by TAO, so there might be unexpected behaviour in overriding this\n",
    "                          \"override_automl_disabled_params\": False,\n",
    "                          \"metric\":metric,\n",
    "                          \"automl_add_hyperparameters\":str(additional_automl_parameters),\n",
    "                          \"automl_remove_hyperparameters\":str(remove_default_automl_parameters)\n",
    "                         }\n",
    "    patched_model = subprocess.getoutput(f\"nvtl {model_name} patch-artifact-metadata --id {experiment_id} --job_type experiment --update_info '{json.dumps(automl_information)}' \")\n",
    "    patched_model = json.loads(patched_model)\n",
    "    print(json.dumps(patched_model, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "#### Provide train specs <a class=\"anchor\" id=\"head-15\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default train model specs\n",
    "train_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --action train --job_type experiment --id {experiment_id}\")\n",
    "train_specs = json.loads(train_specs)\n",
    "print(json.dumps(train_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes for any of the parameters listed in the previous cell as required\n",
    "if model_name == \"action_recognition\":\n",
    "    train_specs[\"model\"][\"model_type\"] = model_type\n",
    "    train_specs[\"model\"][\"input_type\"] = model_input_type\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 20\n",
    "    train_specs[\"train\"][\"num_gpus\"] = 1\n",
    "elif model_name == \"bpnet\":\n",
    "    train_specs[\"num_epoch\"] = 20\n",
    "    train_specs[\"checkpoint_n_epoch\"] = 5\n",
    "    train_specs[\"validation_every_n_epoch\"] = 5\n",
    "    train_specs[\"finetuning_config\"][\"checkpoint_path\"] = None\n",
    "    train_specs[\"gpus\"] = 1\n",
    "elif model_name == \"centerpose\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 20 # Please set it to 140 if you want to run the whole training.\n",
    "    train_specs[\"train\"][\"validation_interval\"] = 10\n",
    "    train_specs[\"train\"][\"checkpoint_interval\"] = 10\n",
    "    train_specs[\"train\"][\"num_gpus\"] = 1\n",
    "    train_specs[\"dataset\"][\"category\"] = testing_categories\n",
    "elif model_name == \"fpenet\":\n",
    "    train_specs[\"num_epoch\"] = 10\n",
    "    train_specs[\"checkpoint_n_epoch\"] = 5\n",
    "    train_specs[\"dataloader\"][\"dataset_info\"][\"root_path\"] = None\n",
    "    train_specs[\"num_keypoints\"] = int(model_type)\n",
    "    train_specs[\"dataloader\"][\"num_keypoints\"] = int(model_type)\n",
    "    train_specs[\"gpus\"] = 1\n",
    "elif model_name == \"lprnet\":\n",
    "    train_specs[\"training_config\"][\"num_epochs\"] = 24\n",
    "    train_specs[\"gpus\"] = 1\n",
    "elif model_name == \"ml_recog\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 30\n",
    "    train_specs[\"train\"][\"gpu_ids\"] = [0]\n",
    "    train_specs[\"train\"][\"checkpoint_interval\"] = 5\n",
    "elif model_name == \"ocdnet\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 30\n",
    "    train_specs[\"train\"][\"checkpoint_interval\"] = 5\n",
    "    train_specs[\"train\"][\"validation_interval\"] = 5\n",
    "    train_specs[\"train\"][\"gpu_id\"] = [0]\n",
    "    train_specs[\"num_gpus\"] = 1\n",
    "elif model_name == \"ocrnet\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 20\n",
    "    train_specs[\"train\"][\"checkpoint_interval\"] = 5\n",
    "    train_specs[\"train\"][\"validation_interval\"] = 5\n",
    "    train_specs[\"train\"][\"num_gpus\"] = 1\n",
    "elif model_name == \"optical_inspection\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 20\n",
    "    train_specs[\"train\"][\"checkpoint_interval\"] = 5\n",
    "    train_specs[\"train\"][\"validation_interval\"] = 5\n",
    "    train_specs[\"train\"][\"gpu_ids\"] = [0]\n",
    "elif model_name == \"pose_classification\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 50\n",
    "    train_specs[\"train\"][\"gpu_ids\"] = [0]\n",
    "    train_specs[\"train\"][\"num_gpus\"] = 1\n",
    "    if model_type == \"nvidia\":\n",
    "        train_specs[\"dataset\"][\"num_classes\"] = 6\n",
    "        train_specs[\"model\"][\"graph_layout\"] = \"nvidia\"\n",
    "    elif model_type == \"kinetics\":\n",
    "        train_specs[\"dataset\"][\"num_classes\"] = 5\n",
    "        train_specs[\"model\"][\"graph_layout\"] = \"openpose\"\n",
    "elif model_name == \"pointpillars\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 80\n",
    "    train_specs[\"gpus\"] = 1\n",
    "elif model_name == \"re_identification\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 120\n",
    "    train_specs[\"train\"][\"gpu_ids\"] = [0]\n",
    "    train_specs[\"train\"][\"num_gpus\"] = 1\n",
    "    train_specs[\"dataset\"][\"num_classes\"] = 100 #The number set in obtain_subset script\n",
    "    train_specs[\"dataset\"][\"num_workers\"] = 4 #Modify the num_workers according to your hardware setup\n",
    "    train_specs[\"dataset\"][\"batch_size\"] = 16 #Modify the batch_size according to your hardware setup\n",
    "elif model_name == \"visual_changenet\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 30 \n",
    "    train_specs[\"train\"][\"checkpoint_interval\"] = 2\n",
    "    train_specs[\"train\"][\"val_interval\"] = 5\n",
    "    train_specs[\"num_gpus\"] = 1\n",
    "    if ds_format == \"visual_changenet_segment\":\n",
    "        train_specs[\"task\"] = 'segment'\n",
    "    elif ds_format == \"visual_changenet_classify\":\n",
    "        train_specs[\"task\"] = 'classify'\n",
    "print(json.dumps(train_specs, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "#### Run train <a class=\"anchor\" id=\"head-16\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "ds_convert_parent = \"\"\n",
    "if model_name in (\"bpnet\", \"ocrnet\"):\n",
    "    val_convert_job_id = job_map[\"eval_convert_\" + model_name]\n",
    "    ds_convert_parent = f\"--parent_job_id {val_convert_job_id}\"\n",
    "elif model_name in (\"bpnet\", \"fpenet\", \"ocrnet\", \"pointpillars\"):\n",
    "    train_convert_job_id = job_map[\"train_convert_\" + model_name]\n",
    "    ds_convert_parent = f\"--parent_job_id {train_convert_job_id}\"\n",
    "\n",
    "job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action train --id {experiment_id} --specs '{json.dumps(train_specs)}' {ds_convert_parent}\")\n",
    "if model_name == \"visual_changenet\":\n",
    "    job_map[\"train_\" + ds_format] = job_id\n",
    "else:\n",
    "    job_map[\"train_\" + model_name] = job_id\n",
    "print(job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Monitor job status\n",
    "if automl_enabled:    \n",
    "    while True:\n",
    "        clear_output(wait=True)\n",
    "        response = subprocess.getoutput(f\"nvtl {model_name} get-action-status --job_type experiment --id {experiment_id} --job {job_id}\")\n",
    "        response = json.loads(response)\n",
    "        if \"error_desc\" in response.keys() and response[\"error_desc\"] in (\"Job trying to retrieve not found\", \"No AutoML run found\"):\n",
    "            print(\"Job is being created\")\n",
    "            time.sleep(5)\n",
    "            continue\n",
    "        print(json.dumps(response, sort_keys=True, indent=4))\n",
    "        assert \"status\" in response.keys() and response.get(\"status\") != \"Error\"\n",
    "        if response.get(\"status\") in [\"Done\",\"Error\"]:\n",
    "            break\n",
    "        time.sleep(15)\n",
    "else:\n",
    "    # Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "    status = my_tail(model_name, experiment_id, job_id, \"experiment\", workdir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## To Stop an AutoML JOB\n",
    "#    1. Stop the 'Monitor job status' cell (the cell right before this cell) manually\n",
    "#    2. Uncomment the snippet in the next cell and run the cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# if automl_enabled:\n",
    "#     if model_name == \"visual_changenet\":\n",
    "#         job_id = job_map[\"train_\" + ds_format]\n",
    "#     else:\n",
    "#         job_id = job_map[\"train_\" + model_name]\n",
    "#     job_id = subprocess.getoutput(f\"nvtl {model_name} job-cancel --job_type experiment --id {experiment_id} --job {job_id}\")\n",
    "#     if model_name == \"visual_changenet\":\n",
    "#         job_map[\"canceled_\" + ds_format] = job_id\n",
    "#     else:\n",
    "#         job_map[\"canceled_\" + model_name] = job_id\n",
    "#     print(job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## Resume AutoML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Uncomment the below snippet if you want to resume an already stopped AutoML job and then run the 'Monitor job status' cell above (4th cell above from this cell)\n",
    "# if automl_enabled:\n",
    "#     if model_name == \"visual_changenet\":\n",
    "#         job_id = job_map[\"train_\" + ds_format]\n",
    "#     else:\n",
    "#         job_id = job_map[\"train_\" + model_name]\n",
    "#     job_id = subprocess.getoutput(f\"nvtl {model_name} job-resume --id {experiment_id} --job {job_id} --specs '{json.dumps(train_specs)}' {ds_convert_parent}\")\n",
    "#     if model_name == \"visual_changenet\":\n",
    "#         job_map[\"resumed_\" + ds_format] = job_id\n",
    "#     else:\n",
    "#         job_map[\"resumed_\" + model_name] = job_id\n",
    "#     print(job_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download train job artifacts <a class=\"anchor\" id=\"head-17\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name == \"visual_changenet\":\n",
    "    job_id = job_map[\"train_\" + ds_format]\n",
    "else:\n",
    "    job_id = job_map[\"train_\" + model_name]\n",
    "file_list = subprocess.getoutput(f\"nvtl {model_name} list-job-files --id {experiment_id} --job {job_id} --job_type experiment --retrieve_logs True --retrieve_specs False\")\n",
    "print(file_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## Patch the model with proper metric before training to run this cell; By default loss is used, but some models dont log the parameter under the name 'loss'\n",
    "# file_lists = []\n",
    "# temptar = subprocess.getoutput(f\"nvtl {model_name} download-selective-files --id {experiment_id} --job {job_id} --job_type experiment --workdir {workdir}  --file_lists '{file_lists}' --best_model False --latest_model True --tar_files True\")\n",
    "# tar_command = f'tar -xvf {temptar} -C {workdir}/'\n",
    "# os.system(tar_command)\n",
    "# os.remove(temptar)\n",
    "# print(f\"Results at {workdir}/{job_id}\")\n",
    "# model_downloaded_path = f\"{workdir}/{job_id}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Downloading train job takes a longer time, uncomment this cell if you want to still proceed\n",
    "if download_jobs:\n",
    "    temptar = subprocess.getoutput(f\"nvtl {model_name} download-entire-job --id {experiment_id} --job {job_id} --job_type experiment --workdir {workdir}\")\n",
    "    tar_command = f'tar -xvf {temptar} -C {workdir}/'\n",
    "    os.system(tar_command)\n",
    "    os.remove(temptar)\n",
    "    print(f\"Results at {workdir}/{job_id}\")\n",
    "    model_downloaded_path = f\"{workdir}/{job_id}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# View the checkpoints generated for the training job and for automl jobs, in addition view: best performing model's config and the results of all automl experiments\n",
    "if download_jobs:\n",
    "    if automl_enabled:\n",
    "        !python3 -m pip install pandas==1.5.1\n",
    "        import pandas as pd\n",
    "        model_downloaded_path = f\"{model_downloaded_path}/best_model\"\n",
    "        assert glob.glob(f\"{model_downloaded_path}/*.protobuf\") or glob.glob(f\"{model_downloaded_path}/*.yaml\")\n",
    "\n",
    "    assert os.path.exists(model_downloaded_path)\n",
    "    assert (glob.glob(model_downloaded_path + \"/**/*.tlt\", recursive=True) + glob.glob(model_downloaded_path + \"/**/*.hdf5\", recursive=True) + glob.glob(model_downloaded_path + \"/**/*.pth\", recursive=True))\n",
    "\n",
    "    if os.path.exists(model_downloaded_path):        \n",
    "        #List the binary model file\n",
    "        print(\"\\nCheckpoints for the training experiment\")\n",
    "        if os.path.exists(model_downloaded_path+\"/train/weights\") and len(os.listdir(model_downloaded_path+\"/train/weights\")) > 0:\n",
    "            print(f\"Folder: {model_downloaded_path}/train/weights\")\n",
    "            print(\"Files:\", os.listdir(model_downloaded_path+\"/train/weights\"))\n",
    "        elif os.path.exists(model_downloaded_path+\"/weights\") and len(os.listdir(model_downloaded_path+\"/weights\")) > 0:\n",
    "            print(f\"Folder: {model_downloaded_path}/weights\")\n",
    "            print(\"Files:\", os.listdir(model_downloaded_path+\"/weights\"))\n",
    "        else:\n",
    "            print(f\"Folder: {model_downloaded_path}\")\n",
    "            print(\"Files:\", os.listdir(model_downloaded_path))\n",
    "\n",
    "        if automl_enabled:\n",
    "            assert glob.glob(f\"{model_downloaded_path}/*.protobuf\") or glob.glob(f\"{model_downloaded_path}/*.yaml\")\n",
    "            experiment_artifacts = json.load(open(f\"{model_downloaded_path}/controller.json\",\"r\"))\n",
    "            data_frame = pd.DataFrame(experiment_artifacts)\n",
    "            # Print experiment id/number and the corresponding result\n",
    "            print(\"\\nResults of all experiments\")\n",
    "            with pd.option_context('display.max_rows', None, 'display.max_columns', None, 'display.max_colwidth', None):\n",
    "                print(data_frame[[\"id\",\"result\"]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate <a class=\"anchor\" id=\"head-16\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Provide evaluate specs <a class=\"anchor\" id=\"head-18\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default evaluate model specs\n",
    "eval_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --action evaluate --job_type experiment --id {experiment_id}\")\n",
    "eval_specs = json.loads(eval_specs)\n",
    "print(json.dumps(eval_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize evaluate model specs\n",
    "if model_name == \"action_recognition\":\n",
    "    eval_specs[\"model\"][\"model_type\"] = model_type\n",
    "    eval_specs[\"model\"][\"input_type\"] = model_input_type\n",
    "elif model_name == \"fpenet\":\n",
    "    eval_specs[\"dataloader\"][\"dataset_info\"][\"root_path\"] = None\n",
    "    eval_specs[\"num_keypoints\"] = int(model_type)\n",
    "    eval_specs[\"dataloader\"][\"num_keypoints\"] = int(model_type)\n",
    "elif model_name == \"pose_classification\":\n",
    "    if model_type == \"nvidia\":\n",
    "        eval_specs[\"dataset\"][\"num_classes\"] = 6\n",
    "        eval_specs[\"model\"][\"graph_layout\"] = \"nvidia\"\n",
    "    elif model_type == \"kinetics\":\n",
    "        eval_specs[\"dataset\"][\"num_classes\"] = 5\n",
    "        eval_specs[\"model\"][\"graph_layout\"] = \"openpose\"\n",
    "elif model_name == \"re_identification\":\n",
    "    eval_specs[\"dataset\"][\"num_classes\"] = 100 #The number set in obtain_subset script\n",
    "elif ds_format == 'visual_changenet_segment':\n",
    "    eval_specs[\"task\"] = 'segment'\n",
    "elif ds_format == 'visual_changenet_classify':\n",
    "    eval_specs[\"task\"] = 'classify'\n",
    "    eval_specs[\"train\"][\"classify\"][\"loss\"] = \"contrastive\"\n",
    "elif model_name == \"centerpose\":\n",
    "    eval_specs[\"dataset\"][\"category\"] = testing_categories\n",
    "print(json.dumps(eval_specs, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run evaluate <a class=\"anchor\" id=\"head-19\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Print model handler parameters\n",
    "model_parameters = subprocess.getoutput(f\"nvtl {model_name} get-metadata --id {experiment_id} --job_type experiment\")\n",
    "model_parameters = json.loads(model_parameters)\n",
    "update_checkpoint_choosing = {}\n",
    "update_checkpoint_choosing[\"checkpoint_choose_method\"] = model_parameters[\"checkpoint_choose_method\"]\n",
    "update_checkpoint_choosing[\"checkpoint_epoch_number\"] = model_parameters[\"checkpoint_epoch_number\"]\n",
    "print(json.dumps(update_checkpoint_choosing, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Change the method by which checkpoint from the parent action is chosen, when parent action is a train/retrain action.\n",
    "# Example for evaluate action below, can be applied in the same way for other actions too\n",
    "update_checkpoint_choosing[\"checkpoint_choose_method\"] = \"latest_model\" # Choose between best_model/latest_model/from_epoch_number\n",
    "# If from_epoch_number is chosen then assign the epoch number to the dictionary key in the format 'from_epoch_number{train_job_id}'\n",
    "# update_checkpoint_choosing[\"checkpoint_epoch_number\"][\"from_epoch_number_c2f76eb7-2a75-4197-9a84-c1547f20c17d\"] = 6\n",
    "\n",
    "patched_model = subprocess.getoutput(f\"nvtl {model_name} patch-artifact-metadata --id {experiment_id} --job_type experiment --update_info '{json.dumps(update_checkpoint_choosing)}' \")\n",
    "patched_model = json.loads(patched_model)\n",
    "print(json.dumps(patched_model, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name == \"visual_changenet\":\n",
    "    parent = job_map[\"train_\" + ds_format]\n",
    "else:\n",
    "    parent = job_map[\"train_\" + model_name]\n",
    "job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action evaluate --id {experiment_id} --parent_job_id {parent} --specs '{json.dumps(eval_specs)}'\")\n",
    "if model_name == \"visual_changenet\":\n",
    "    job_map[\"eval_\" + ds_format] = job_id\n",
    "else:\n",
    "    job_map[\"eval_\" + model_name] = job_id\n",
    "print(job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "status = my_tail(model_name, experiment_id, job_id, \"experiment\", workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prune, Retrain and Evaluation <a class=\"anchor\" id=\"head-13\"></a>\n",
    "\n",
    "- We optimize the trained model by pruning and retraining in the following cells"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "#### Prune <a class=\"anchor\" id=\"head-20\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default prune model specs\n",
    "if model_name in (\"bpnet\", \"ocdnet\", \"ocrnet\", \"pointpillars\"):\n",
    "    prune_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --id {experiment_id} --action prune --job_type experiment\")\n",
    "    prune_specs = json.loads(prune_specs)\n",
    "    print(json.dumps(prune_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply changes to prune specs if neccessary\n",
    "if model_name in (\"bpnet\", \"ocdnet\", \"ocrnet\", \"pointpillars\"):\n",
    "    print(json.dumps(prune_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if model_name in (\"bpnet\",\"pointpillars\"):\n",
    "    parent = job_map[\"train_\" + model_name]\n",
    "    job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action prune --id {experiment_id} --parent_job_id {parent} --specs '{json.dumps(prune_specs)}'\")\n",
    "    job_map[\"prune_\" + model_name] = job_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check status of pruning job (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "if model_name in (\"bpnet\",\"pointpillars\"):\n",
    "    status = my_tail(model_name, experiment_id, job_id, \"experiment\", workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "#### Retrain <a class=\"anchor\" id=\"head-22\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default retrain model specs\n",
    "if model_name in (\"bpnet\", \"ocdnet\", \"ocrnet\", \"pointpillars\"):\n",
    "    retrain_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --id {experiment_id} --action retrain --job_type experiment\")\n",
    "    retrain_specs = json.loads(retrain_specs)\n",
    "    print(json.dumps(retrain_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize retrain model specs\n",
    "if model_name in (\"bpnet\", \"ocdnet\", \"ocrnet\", \"pointpillars\"):\n",
    "    if model_name == \"bpnet\":\n",
    "        retrain_specs[\"finetuning_config\"][\"checkpoint_path\"] = None\n",
    "        retrain_specs[\"gpus\"] = 1\n",
    "    elif model_name == \"ocdnet\":\n",
    "        retrain_specs[\"train\"][\"num_epochs\"] = 30\n",
    "        retrain_specs[\"train\"][\"gpu_id\"] = [0]\n",
    "        retrain_specs[\"num_gpus\"] = 1\n",
    "    elif model_name == \"ocrnet\":\n",
    "        retrain_specs[\"train\"][\"num_epochs\"] = 20\n",
    "        retrain_specs[\"train\"][\"gpu_ids\"] = [0]\n",
    "    elif model_name == \"pointpillars\":\n",
    "        retrain_specs[\"train\"][\"num_epochs\"] = 80\n",
    "        retrain_specs[\"gpus\"] = 1\n",
    "    print(json.dumps(retrain_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if model_name in (\"bpnet\",\"pointpillars\"):\n",
    "    parent = job_map[\"prune_\" + model_name]\n",
    "    job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action retrain --id {experiment_id} --parent_job_id {parent} --specs '{json.dumps(retrain_specs)}'\")\n",
    "    job_map[\"retrain_\" + model_name] = job_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check status of retrain job (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "if model_name in (\"bpnet\",\"pointpillars\"):\n",
    "    status = my_tail(model_name, experiment_id, job_id, \"experiment\", workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Evaluate after retrain <a class=\"anchor\" id=\"head-18\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Default evaluate model specs\n",
    "if model_name in (\"bpnet\", \"ocdnet\", \"ocrnet\", \"pointpillars\"):\n",
    "    eval_retrain_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --action evaluate --job_type experiment --id {experiment_id}\")\n",
    "    eval_retrain_specs = json.loads(eval_retrain_specs)\n",
    "    print(json.dumps(eval_retrain_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Customize evaluate model specs if necessary\n",
    "if model_name in (\"bpnet\", \"ocdnet\", \"ocrnet\", \"pointpillars\"):\n",
    "    print(json.dumps(eval_retrain_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if model_name in (\"bpnet\",\"pointpillars\"):\n",
    "    parent = job_map[\"retrain_\" + model_name]\n",
    "    job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action evaluate --id {experiment_id} --parent_job_id {parent} --specs '{json.dumps(eval_retrain_specs)}'\")\n",
    "    job_map[\"eval2_\" + model_name] = job_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "if model_name in (\"bpnet\",\"pointpillars\"):\n",
    "    status = my_tail(model_name, experiment_id, job_id, \"experiment\", workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Export <a class=\"anchor\" id=\"head-22\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Provide Export specs <a class=\"anchor\" id=\"head-25\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default export model specs\n",
    "export_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --id {experiment_id} --action export --job_type experiment\")\n",
    "export_specs = json.loads(export_specs)\n",
    "print(json.dumps(export_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize export model specs\n",
    "if model_name == \"action_recognition\":\n",
    "    export_specs[\"model\"][\"model_type\"] = model_type\n",
    "    export_specs[\"model\"][\"input_type\"] = model_input_type\n",
    "elif model_name == \"lprnet\":\n",
    "    export_specs[\"data_type\"] = \"fp32\"\n",
    "elif model_name == \"bpnet\":\n",
    "    export_specs[\"data_type\"] = \"int8\"\n",
    "elif model_name == \"pose_classification\":\n",
    "    if model_type == \"nvidia\":\n",
    "        export_specs[\"dataset\"][\"num_classes\"] = 6\n",
    "        export_specs[\"model\"][\"graph_layout\"] = \"nvidia\"\n",
    "    elif model_type == \"kinetics\":\n",
    "        export_specs[\"dataset\"][\"num_classes\"] = 5\n",
    "        export_specs[\"model\"][\"graph_layout\"] = \"openpose\"\n",
    "elif model_name == \"re_identification\":\n",
    "    export_specs[\"dataset\"][\"num_classes\"] = 100 #The number set in obtain_subset script\n",
    "elif ds_format == 'visual_changenet_segment':\n",
    "    export_specs[\"export\"][\"input_height\"] = 256 \n",
    "    export_specs[\"export\"][\"input_width\"] = 256 \n",
    "    export_specs[\"task\"] = 'segment'\n",
    "elif ds_format == 'visual_changenet_classify':\n",
    "    export_specs[\"export\"][\"input_height\"] = 512 \n",
    "    export_specs[\"export\"][\"input_width\"] = 128\n",
    "    export_specs[\"task\"] = 'classify'\n",
    "print(json.dumps(export_specs, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run export <a class=\"anchor\" id=\"head-26\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name == \"visual_changenet\":\n",
    "    parent = job_map[\"train_\" + ds_format]\n",
    "else:\n",
    "    parent = job_map[\"train_\" + model_name]\n",
    "job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action export --id {experiment_id} --parent_job_id {parent} --specs '{json.dumps(export_specs)}'\")\n",
    "if model_name == \"visual_changenet\":\n",
    "    job_map[\"export_\" + ds_format] = job_id\n",
    "else:\n",
    "    job_map[\"export_\" + model_name] = job_id\n",
    "print(job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "status = my_tail(model_name, experiment_id, job_id, \"experiment\", workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TRT Engine generation using TAO-Deploy <a class=\"anchor\" id=\"head-26\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Provide trt engine generation specs <a class=\"anchor\" id=\"head-27\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default gen_trt_engine model specs\n",
    "if model_name in (\"bpnet\", \"lprnet\", \"ocdnet\", \"ocrnet\", \"ml_recog\",\"optical_inspection\", \"visual_changenet\", \"centerpose\"):\n",
    "    if model_name in (\"lprnet\", \"ocdnet\", \"ocrnet\", \"ml_recog\",\"optical_inspection\", \"visual_changenet\", \"centerpose\"):\n",
    "        tao_deploy_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --id {experiment_id} --action gen_trt_engine --job_type experiment\")\n",
    "    elif model_name in (\"bpnet\"):\n",
    "        tao_deploy_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --id {experiment_id} --action trtexec --job_type experiment\")\n",
    "    tao_deploy_specs = json.loads(tao_deploy_specs)\n",
    "    print(json.dumps(tao_deploy_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize convert model specs\n",
    "if model_name in (\"bpnet\", \"lprnet\", \"ocdnet\", \"ocrnet\", \"ml_recog\",\"optical_inspection\", \"visual_changenet\"):\n",
    "    # Make changes to the specs dictionary if required here\n",
    "    if model_name == \"lprnet\":\n",
    "        tao_deploy_specs[\"data_type\"] = \"fp32\"\n",
    "    elif model_name in (\"ml_recog\", \"ocdnet\"):\n",
    "        tao_deploy_specs[\"gen_trt_engine\"][\"tensorrt\"][\"data_type\"] = \"int8\"\n",
    "    elif model_name in (\"ocrnet\", \"optical_inspection\"):\n",
    "        tao_deploy_specs[\"gen_trt_engine\"][\"tensorrt\"][\"data_type\"] = \"fp16\"   \n",
    "    elif ds_format == 'visual_changenet_classify':\n",
    "        tao_deploy_specs[\"gen_trt_engine\"][\"input_height\"] = 512 \n",
    "        tao_deploy_specs[\"gen_trt_engine\"][\"input_width\"] = 128\n",
    "        tao_deploy_specs[\"task\"] = 'classify'\n",
    "    elif ds_format == 'visual_changenet_segment':\n",
    "        tao_deploy_specs[\"gen_trt_engine\"][\"tensorrt\"][\"data_type\"] = \"fp16\"\n",
    "        tao_deploy_specs[\"gen_trt_engine\"][\"input_height\"] = 256\n",
    "        tao_deploy_specs[\"gen_trt_engine\"][\"input_width\"]= 256\n",
    "        tao_deploy_specs[\"task\"] = 'segment' \n",
    "    print(json.dumps(tao_deploy_specs, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run TRT engine generation using TAO-Deploy <a class=\"anchor\" id=\"head-28\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name in (\"bpnet\", \"lprnet\", \"ocdnet\", \"ocrnet\", \"ml_recog\",\"optical_inspection\", \"visual_changenet\", \"centerpose\"):\n",
    "    if model_name == \"visual_changenet\":\n",
    "        parent = job_map[\"export_\" + ds_format]\n",
    "    else:\n",
    "        parent = job_map[\"export_\" + model_name]\n",
    "    if model_name in (\"lprnet\", \"ocdnet\", \"ocrnet\", \"ml_recog\",\"optical_inspection\", \"visual_changenet\", \"centerpose\"):\n",
    "        job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action gen_trt_engine --id {experiment_id} --parent_job_id {parent} --specs '{json.dumps(tao_deploy_specs)}'\")\n",
    "    elif model_name in (\"bpnet\"):\n",
    "        job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action trtexec --id {experiment_id} --parent_job_id {parent} --specs '{json.dumps(tao_deploy_specs)}'\")\n",
    "    if model_name == \"visual_changenet\":\n",
    "        job_map[\"gen_trt_engine_\" + ds_format] = job_id\n",
    "    else:\n",
    "        job_map[\"gen_trt_engine_\" + model_name] = job_id\n",
    "    print(job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "if model_name in (\"bpnet\", \"lprnet\", \"ocdnet\", \"ocrnet\", \"ml_recog\",\"optical_inspection\", \"visual_changenet\", \"centerpose\"):\n",
    "    status = my_tail(model_name, experiment_id, job_id, \"experiment\", workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TAO inference <a class=\"anchor\" id=\"head-28\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Provide TAO inference specs <a class=\"anchor\" id=\"head-29\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default inference model specs\n",
    "tao_inference_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --id {experiment_id} --action inference --job_type experiment\")\n",
    "tao_inference_specs = json.loads(tao_inference_specs)\n",
    "print(json.dumps(tao_inference_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize TAO inference specs\n",
    "if model_name == \"action_recognition\":\n",
    "    tao_inference_specs[\"model\"][\"model_type\"] = model_type\n",
    "    tao_inference_specs[\"model\"][\"input_type\"] = model_input_type\n",
    "elif model_name == \"fpenet\":\n",
    "    tao_inference_specs[\"num_keypoints\"] = int(model_type)\n",
    "    tao_inference_specs[\"dataloader\"][\"num_keypoints\"] = int(model_type)\n",
    "elif model_name == \"pose_classification\":\n",
    "    if model_type == \"nvidia\":\n",
    "        tao_inference_specs[\"dataset\"][\"num_classes\"] = 6\n",
    "        tao_inference_specs[\"model\"][\"graph_layout\"] = \"nvidia\"\n",
    "    elif model_type == \"kinetics\":\n",
    "        tao_inference_specs[\"dataset\"][\"num_classes\"] = 5\n",
    "        tao_inference_specs[\"model\"][\"graph_layout\"] = \"openpose\"\n",
    "elif model_name == \"re_identification\":\n",
    "    tao_inference_specs[\"dataset\"][\"num_classes\"] = 100 #The number set in obtain_subset script\n",
    "elif ds_format == 'visual_changenet_classify':\n",
    "    tao_inference_specs[\"inference\"][\"batch_size\"] = tao_inference_specs[\"dataset\"][\"classify\"]['batch_size'] \n",
    "    tao_inference_specs[\"task\"] = 'classify'\n",
    "elif ds_format == 'visual_changenet_segment':\n",
    "    tao_inference_specs[\"inference\"][\"batch_size\"] = tao_inference_specs[\"dataset\"][\"segment\"]['batch_size'] \n",
    "    tao_inference_specs[\"task\"] = 'segment'\n",
    "elif model_name == \"centerpose\":\n",
    "    tao_inference_specs[\"dataset\"][\"category\"] = testing_categories\n",
    "print(json.dumps(tao_inference_specs, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run TAO inference <a class=\"anchor\" id=\"head-30\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name == \"visual_changenet\":\n",
    "    parent = job_map[\"train_\" + ds_format]\n",
    "else:\n",
    "    parent = job_map[\"train_\" + model_name]\n",
    "job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action inference --id {experiment_id} --parent_job_id {parent} --specs '{json.dumps(tao_inference_specs)}'\")\n",
    "if model_name == \"visual_changenet\":\n",
    "    job_map[\"tlt_inference_\" + ds_format] = job_id\n",
    "else:\n",
    "    job_map[\"tlt_inference_\" + model_name] = job_id\n",
    "print(job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "status = my_tail(model_name, experiment_id, job_id, \"experiment\", workdir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if download_jobs:\n",
    "    temptar = subprocess.getoutput(f\"nvtl {model_name} download-entire-job --id {experiment_id} --job {job_id} --job_type experiment --workdir {workdir}\")\n",
    "    tar_command = f'tar -xvf {temptar} -C {workdir}/'\n",
    "    os.system(tar_command)\n",
    "    os.remove(temptar)\n",
    "    print(f\"Results at {workdir}/{job_id}\")\n",
    "    inference_out_path = f\"{workdir}/{job_id}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Inference output must be here\n",
    "if download_jobs:\n",
    "    if model_name in (\"action_recognition\",\"lprnet\",\"ocrnet\"):\n",
    "        assert os.path.exists(f'{inference_out_path}/logs_from_toolkit.txt')\n",
    "        !cat {inference_out_path}/logs_from_toolkit.txt\n",
    "    elif model_name in (\"bpnet\",\"pointpillars\"):\n",
    "        if model_name == \"bpnet\":\n",
    "            assert glob.glob(f\"{inference_out_path}/images_annotated/*.png\")\n",
    "        elif model_name == \"pointpillars\":\n",
    "            assert glob.glob(f\"{inference_out_path}/infer/detected_boxes/*.png\")\n",
    "        elif model_name == \"centerpose\":\n",
    "            assert glob.glob(f\"{inference_out_path}/inference/*.png\")\n",
    "        !python3 -m pip install matplotlib\n",
    "        import matplotlib.pyplot as plt\n",
    "        import matplotlib.image as mpimg\n",
    "        if model_name == \"bpnet\":\n",
    "            sample_image = glob.glob(f\"{inference_out_path}/images_annotated/*.png\")[0]\n",
    "        elif model_name == \"pointpillars\":\n",
    "            sample_image = glob.glob(f\"{inference_out_path}/infer/detected_boxes/*.png\")[0]\n",
    "        elif model_name == \"centerpose\":\n",
    "            sample_image = glob.glob(f\"{inference_out_path}/inference/*.png\")[0]\n",
    "        def display_photo(path):\n",
    "            img = mpimg.imread(path)\n",
    "            plt.figure(figsize = (int(img.shape[0]/100)*2,int(img.shape[1]/100)*2))\n",
    "            plt.axis('off')\n",
    "            imgplot = plt.imshow(img, aspect='auto')\n",
    "            plt.show()\n",
    "        display_photo(sample_image)\n",
    "    elif model_name == \"fpenet\":\n",
    "        assert os.path.exists(f'{inference_out_path}/result.txt')\n",
    "        !cat {inference_out_path}/result.txt\n",
    "    elif model_name == \"ml_recog\":\n",
    "        assert os.path.exists(f'{inference_out_path}/inference/result.csv')\n",
    "        !cat {inference_out_path}/inference/result.csv\n",
    "    elif model_name == \"optical_inspection\":\n",
    "        assert os.path.exists(f'{inference_out_path}/inference/inference.csv')\n",
    "        !cat {inference_out_path}/inference/inference.csv\n",
    "    elif model_name == \"pose_classification\":\n",
    "        assert os.path.exists(f'{inference_out_path}/results.txt')\n",
    "        !cat {inference_out_path}/results.txt\n",
    "    elif model_name == \"re_identification\":\n",
    "        assert os.path.exists(f'{inference_out_path}/inference.json')\n",
    "        !cat {inference_out_path}/inference.json\n",
    "    elif model_name == \"visual_changenet\":\n",
    "        if ds_format == 'visual_changenet_classify':\n",
    "            assert os.path.exists(f'{inference_out_path}/inference/inference.csv')\n",
    "            !cat {inference_out_path}/inference/inference.csv\n",
    "        elif ds_format == 'visual_changenet_segment':\n",
    "            assert os.path.exists(f'{inference_out_path}/inference/status.json')\n",
    "            !cat {inference_out_path}/inference/status.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TRT inference <a class=\"anchor\" id=\"head-30\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Provide TRT inference specs <a class=\"anchor\" id=\"head-31\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default inference model specs\n",
    "if model_name in (\"bpnet\", \"lprnet\", \"ocdnet\", \"ocrnet\", \"ml_recog\",\"optical_inspection\", \"visual_changenet\", \"centerpose\"):\n",
    "    trt_inference_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --id {experiment_id} --action inference --job_type experiment\")\n",
    "    trt_inference_specs = json.loads(trt_inference_specs)\n",
    "    print(json.dumps(trt_inference_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize TRT inference specs\n",
    "# Change any spec if you wish\n",
    "if model_name in (\"bpnet\", \"lprnet\", \"ocdnet\", \"ocrnet\", \"ml_recog\",\"optical_inspection\", \"visual_changenet\", \"centerpose\"):\n",
    "    if ds_format == 'visual_changenet_classify':\n",
    "        trt_inference_specs[\"inference\"][\"batch_size\"] = trt_inference_specs[\"dataset\"][\"classify\"]['batch_size']\n",
    "        trt_inference_specs[\"task\"] = 'classify'\n",
    "    elif ds_format == 'visual_changenet_segment':\n",
    "        trt_inference_specs[\"inference\"][\"batch_size\"] = trt_inference_specs[\"dataset\"][\"segment\"]['batch_size']\n",
    "        trt_inference_specs[\"task\"] = 'segment'\n",
    "    print(json.dumps(trt_inference_specs, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run TRT inference <a class=\"anchor\" id=\"head-32\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name in (\"bpnet\", \"lprnet\", \"ocdnet\", \"ocrnet\", \"ml_recog\",\"optical_inspection\", \"visual_changenet\", \"centerpose\"):\n",
    "    if model_name == \"visual_changenet\":\n",
    "        parent = job_map[\"gen_trt_engine_\" + ds_format]\n",
    "    else:\n",
    "        parent = job_map[\"gen_trt_engine_\" + model_name]\n",
    "    job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action inference --id {experiment_id} --parent_job_id {parent} --specs '{json.dumps(trt_inference_specs)}'\")\n",
    "    if model_name == \"visual_changenet\":\n",
    "        job_map[\"trt_inference_\" + ds_format] = job_id\n",
    "    else:\n",
    "        job_map[\"trt_inference_\" + model_name] = job_id\n",
    "    print(job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "if model_name in (\"bpnet\", \"lprnet\", \"ml_recog\", \"ocdnet\", \"ocrnet\",\"optical_inspection\", \"visual_changenet\", \"centerpose\"):\n",
    "    status = my_tail(model_name, experiment_id, job_id, \"experiment\", workdir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if download_jobs:\n",
    "    if model_name in (\"bpnet\", \"lprnet\", \"ml_recog\", \"ocdnet\", \"ocrnet\",\"optical_inspection\", \"visual_changenet\", \"centerpose\"):\n",
    "        temptar = subprocess.getoutput(f\"nvtl {model_name} download-entire-job --id {experiment_id} --job {job_id} --job_type experiment --workdir {workdir}\")\n",
    "        tar_command = f'tar -xvf {temptar} -C {workdir}/'\n",
    "        os.system(tar_command)\n",
    "        os.remove(temptar)\n",
    "        print(f\"Results at {workdir}/{job_id}\")\n",
    "        job_dir = f\"{workdir}/{job_id}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if download_jobs:\n",
    "    if model_name in (\"bpnet\", \"lprnet\", \"ml_recog\", \"ocdnet\", \"ocrnet\",\"optical_inspection\", \"visual_changenet\", \"centerpose\"):\n",
    "        # You can view the predictions here or in the subdirectories\n",
    "        !ls {job_dir}/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Delete model <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subprocess.getoutput(f\"nvtl {model_name} experiment-delete --id {experiment_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Delete dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Delete train dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subprocess.getoutput(f\"nvtl {model_name} dataset-delete --id {train_dataset_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Delete val dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if model_name in (\"lprnet\",\"ocdnet\",\"ocrnet\",\"optical_inspection\") or ds_format == 'visual_changenet_classify':\n",
    "    subprocess.getoutput(f\"nvtl {model_name} dataset-delete --id {eval_dataset_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Delete test dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if model_name in (\"lprnet\", \"optical_inspection\") or ds_format == 'visual_changenet_classify':\n",
    "    subprocess.getoutput(f\"nvtl {model_name} dataset-delete --id {test_dataset_id}\")"
   ]
  }
 ],
 "metadata": {
  "datalore": {
   "base_environment": "default",
   "computation_mode": "JUPYTER",
   "package_manager": "pip",
   "packages": [],
   "version": 1
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
